import taichi as ti
import random

ti.init()

n=8

x = ti.field(dtype=ti.f32,shape=n,needs_grad=True)
y = ti.field(dtype=ti.f32,shape=n)
L = ti.field(dtype=ti.f32,shape=(),needs_grad=True)


## loss函数：0.5 * (x-y)**2
@ti.kernel
def reduce():
    for i in range(n):
        L[None] += 0.5 * (x[i]-y[i])**2

# 随机生成x,y
for i in range(n):
    x[i] = random.random()
    y[i] = random.random()


# 梯度下降
@ti.kernel
def gradient_descend():
    for i in x:
        x[i] -= x.grad[i] *0.1

# 200轮
for k in range(200):
    # 自动计算loss
    with ti.Tape(loss=L):
        reduce()
    # 打印loss
    print('loss={}'.format(L[None]))
    # 用梯度下降调整x的值
    gradient_descend()

for i in range(n):
    print(x[i],y[i])